A Soware-defined Tensor Streaming Multiprocessor for Large-scale Machine Learning

https://scrapbox.io/files/6628a68ee2242a0023855f3d.png

論文情報

タイトル：A Soware-defined Tensor Streaming Multiprocessor for Large-scale Machine Learning

発行日：2022年6月

著者：Dennis Abts etal

所属：Groq

論文のポイント

この論文は、たくさんの機械学習の計算を高速に行うための新しいコンピュータシステムについて説明している。

このシステムでは、「TSP」と呼ばれる特別なプロセッサを使っている。

TSPは、計算の手順が予め決められているので、とても効率的に動作する。

たくさんのTSPを協力させて働かせるために、「ドラゴンフライ」と呼ばれるネットワークの形を使っている。

これにより、最大で10,440個ものTSPを接続できる。

TSPは、全体で2テラバイト以上のメモリを持っている。

どのTSPからでも、このメモリに3マイクロ秒（1秒の300万分の1）以内にアクセス可能。

コンピュータが行う計算と通信は、あらかじめソフトウェアが計画を立てる。

これによって、無駄なく効率的に動作させることができる。

行列の計算、自然言語処理、データの集計など、様々な種類の計算を高速に処理できることを確認している。

このシステムには、PyTorchやTensorFlowといった機械学習のためのソフトウェアを使うこともできる。

ソフトウェアとハードウェアが協力することで、高速で無駄のない計算を実現していることがポイント。

https://scrapbox.io/files/6628aa147c1bfb0024ea91ec.png

Figure 1は、従来のマルチコアチップマルチプロセッサの概略図。

図中で赤色で強調されている部分は、メモリやネットワークの参照を並べ替える可能性のある非決定性の源。

中央には複数のプロセッシングコア（Processing Cores）がある。

各コアには、命令キャッシュ（I-Cache）とデータキャッシュ（D-Cache）が付属している。

これらのキャッシュは、メインメモリへのアクセスを高速化するために使用される。

各コアはインターコネクトスイッチ（Interconnect Switch）に接続されており、これによってコア間の通信が可能になる。

また、インターコネクトスイッチは、共有されたL2キャッシュ（Shared L2 Cache）とも接続されている。

メモリコントローラ（Memory Controller）は、メインメモリ（Main Memory）へのアクセスを管理する。

DMAエンジン（DMA Engine）は、メモリとI/Oデバイス間のデータ転送を担当する。

これに対し、本論文で提案されているTSPアーキテクチャでは、このような非決定性を排除するために、ソフトウェアによるスケジューリングを採用している。

https://scrapbox.io/files/6628aaae827f9b002ddca0dd.png

Figure 2は、システムの規模（TSPの数）に対する、TSP1個あたりの利用可能なグローバル帯域幅の関係を示したグラフ。

つまり、システムのサイズが大きくなるにつれて、各TSPが利用できるグローバル帯域幅がどのように変化するかを表している。

グラフの横軸はシステムの規模（エンドポイントの数）、縦軸はTSP1個あたりのグローバル帯域幅（GB/s）を表している。

特徴的な点は以下の通りです：

1. 16個未満のTSPを持つ小規模システムでは、ノード内の豊富な配線を活用できるため、TSPあたりの帯域幅が非常に高くなっている。

2. 数百個のTSPまでの大規模システムでは、TSPあたり約50GB/sのグローバル（バイセクション）帯域幅を利用できる。

3. システムサイズが264個のTSPを超えると、利用可能なグローバル帯域幅はTSPあたり約14GB/sに下がる。

4. Dragonfly ネットワークは、最大構成の145キャビネット（合計10,440個のTSP）まで、一定のグローバル帯域幅を提供する。

小規模なシステムではTSPあたりの帯域幅が非常に高く、システムが大規模になるにつれて帯域幅が減少していくが、Dragonflyネットワークを使うことで、大規模なシステムでも一定の帯域幅を維持できることがわかる。

https://scrapbox.io/files/6628ab3abc2c450025699ee7.png

Figure 3は、グローバル共有アドレス空間がシステム内のTSPに物理的に分散されている様子を示した図。

図の上部には、連続したアドレス空間が示されている。

これは、ソフトウェアから見えるグローバルメモリの論理的なアドレス空間。

アプリケーションやプログラムは、このアドレス空間を連続的で一体的なものとして扱うことができる。

しかし、実際のハードウェアでは、このグローバルメモリは物理的に分散されている。

図の下部には、複数のTSPが並んでいる。

各TSPには、オンチップのSRAMが搭載されており、これがグローバルメモリの一部を構成している。

つまり、論理的には一つの大きな共有メモリに見えるグローバルメモリが、物理的には複数のTSPのローカルメモリに分散して存在している。

一方、ハードウェアは、アクセスされたアドレスに応じて、適切なTSPのローカルメモリにアクセスする。

このように、グローバル共有アドレス空間の概念を導入することで、プログラミングモデルの簡略化と、物理的なメモリの分散を両立している。これは、大規模なシステムを構築する上で重要な技術となる。

https://scrapbox.io/files/6628abf83871e80023c3f219.png

Figure 4は、間接ネットワーク、直接ネットワーク、およびTSPシステムで提案されているソフトウェアスケジュールドネットワークの構成を比較している。

(a) 間接ネットワーク（Indirect Network）：

処理エンドポイント（P）とルータ（R）が別々のノードとして存在する。処理エンドポイントは、ルータを介して他の処理エンドポイントと通信する。この構成では、ルータが通信の中継点となるため、ルータでの輻輳やオーバーヘッドが発生する可能性がある。

(b) 直接ネットワーク（Direct Network）：

処理エンドポイント（P）とルータ（R）が同一のノード内に存在します。各ノードは直接他のノードと通信可能ですが、依然としてルータが通信を管理している。この構成では、間接ネットワークと比較して、通信のオーバーヘッドが減少する。

(c) ソフトウェアスケジュールドネットワーク（Software-Scheduled Network）（TSPシステムで提案）：

処理エンドポイント（P）とルータ（R）の機能が完全に統合されている。各ノードは処理とルーティングの両方の役割を担う。この構成では、ルータがなくなることで、ハードウェアの複雑さが減少し、通信のオーバーヘッドがさらに減少する。また、ソフトウェアによるスケジューリングが可能となり、通信の最適化や決定論的な動作が実現できる。

https://scrapbox.io/files/6628ac639336d0002477d5e7.png

Figure 5は、システムの4つのパッケージ階層を示している。

(a) チップ（Chip）：

TSPチップ自体を表している。各チップにはヒートシンクが取り付けられている。

(b) PCIeカード（PCIe Card）：

TSPチップはPCIeカードに搭載されている。このカードがシステムの基本的な構成要素となる。

(c) シェアードメモリノード（Shared Memory Node）：

4Uの筐体（シャーシ）内に8枚のTSPカードが搭載されている。この8枚のカードで1つのノードを構成する。ノード内では、TSPカード間を低背ケーブルで接続することで、高速な通信を実現していル。

(d) ラック（Rack）：

複数のノードを格納するための筐体。1つのノードは4Uの高さを占める。

図の右側には、1つのノードの拡大図が示されている。

ノード内では、8枚のTSPカードが密に接続されており、これによって、ノード内の通信を高速化している。

このようなパッケージ階層を採用することで、システムの物理的な構成を最適化している。チップ、カード、ノード、ラックと、階層が上がるにつれて、搭載できるTSPの数が増加する。また、下位の階層ほど高速な接続が可能となり、上位の階層ほど多数のTSPを収容できるようになる。